'''
Created on Aug 4, 2014

@author: Siyuan Qi
'''

import sys
import os
import xml.etree.ElementTree as ET

# Dependencies
import numpy as np
import pydot
import pymc as pm
from pymc.distributions import rcategorical

# Local imports
from grammarNode import Terminal, SetProduction
from entity import DistParam, Entity
from sputil import smallNumber, calCollisionVolume


# ====================== Class of Grammar Tree ======================
class GrammarModel:
	def __init__(self, xml, cameraHeight, saveGraph, folder):
		""" Initializes the grammar model

		Read in the XML grammar file, and generate the scene grammar tree.  The root node of the 
		grammar tree is self.grammarTreeRoot, and the root node of the parse tree is self.parseTreeRoot.

		Args:
			xml: Relative path to the xml file
			cameraHeight: The camera height in the scene.  Since the coordinates in the proposed 
				cuboids are based on the coordinate system using the camera as the origin, the mean 
				position of 'z' value of each object will be deducted by cameraHeight.
		"""

		# Functions that expands the grammar tree
		def parseNode(node, nodes, productions, distributions):
			if node.text not in distributions:
				print('Error in ParseNode: Unspecified distribution of "{:s}"'.format(node.text))
				sys.exit(1)

			if node.tag == 'Root' or node.tag == 'Nonterminal' :
				children = []
				childrenProb = {}
				# First created the children nodes
				for child in productions[node.text]:
					c = parseNode(nodes[child.attrib['symbol']], nodes, productions, distributions)
					children.append(c)
					childrenProb[child.attrib['symbol']] = [float(p) for p in child.attrib['prob'].split()]

				# Create the parent node
				parentNode = SetProduction(node, children, childrenProb, distributions[node.text], cameraHeight)

				# Assign children to the parent node
				for child in children:
					child.parent = parentNode

				return parentNode
			elif node.tag == 'Terminal' :
				return Terminal(node, distributions[node.text], cameraHeight)
			else:
				print('ERROR in ParseNode: Could not locate child name %s', node.text)
				sys.exit(1)

		def expandXMLRoot(xmlroot):
			nodes = {}
			for t in xmlroot.find('Symbols'):
				nodes[t.text] = t

			productions = {}
			for t in xmlroot.find('Productions').findall('SET'):
				productions[t.attrib['symbol']] = t

			distributions = {}
			for t in xmlroot.find('Distributions').findall('Parameters'):
				distributions[t.attrib['symbol']] = DistParam(t)

			# Return the root of the tree structure
			return parseNode(xmlroot.find('Symbols').find('Root'), nodes, productions, distributions)

		def getParsedGrammarRoot(xml):
			parsedxml = ET.parse(xml)
			xmlroot = parsedxml.getroot()
			return expandXMLRoot(xmlroot)


		self.xml = xml
		self.grammarTreeRoot = getParsedGrammarRoot(xml)
		self.parseTreeRoot = None
		
		self.T = 1.
		
		self.saveGraph = saveGraph
		self.folder = folder

	def initializeParseTree(self, bgDist, cubDist):
		""" Initializes a parse tree.

		Top down sampling a parse tree from the root node.

		Args:
			bg: The proposed cuboid of the background.  If there is a proposal of the background, 
				the cuboid will be assigned to the root entity, and sample the children entities
				afterwards.  If there is no proposal, a parse tree will be sampled directly according
				to the defined gramamr.
		"""
		def initializeFromProposal(entity):
			if entity != self.parseTreeRoot and entity.node.isEntity:
				entity.assignCuboid(self.cubDist.propose(entity.node.symbol))
			
			if entity.children:
				for child in entity.children:
					initializeFromProposal(child)
				
		rootEntity = Entity(self.grammarTreeRoot, None)
		rootEntity.sampleAttribute()
		self.bgDist = bgDist
		self.cubDist = cubDist
		rootEntity.assignCuboid(self.bgDist.getDataProposal())
		self.parseTreeRoot = self.grammarTreeRoot.sample(rootEntity)
		self.drawGraph('initialGraph.png')
		
		#initializeFromProposal(self.parseTreeRoot)

	def randomTest(self):
		self.printTree(self.parseTreeRoot, 0)
		count = {}
		for i in range(10000):
			rnode = self.parseTreeRoot.randomTermEntity()
			if rnode.node.symbol in count:
				count[rnode.node.symbol] += 1
			else:
				count[rnode.node.symbol] = 1

		print 'Total number of nodes in the tree:', self.parseTreeRoot.termEntitiesCount()
		print 'Result:', count

	def diffusion(self):
		def proposalLogProb(chosenEntity):
			totalProb = 0
			
			if chosenEntity == self.parseTreeRoot:
				# alpha channel
				prob = self.bgDist.evaluate(chosenEntity.getCuboid())
				prob += smallNumber
				totalProb += channelProb[0]*(1./self.parseTreeRoot.termEntitiesCount())*prob
				
				# beta channel
				if chosenEntity.children:
					prob = np.exp(chosenEntity.proposeFromChildLogProb())
					prob += smallNumber
					totalProb += channelProb[1]*(1./self.parseTreeRoot.termEntitiesCount())*prob
				
				# gamma channel
				prob = np.exp(chosenEntity.calAttLogProb())
				prob += smallNumber
				totalProb += channelProb[2]*(1./self.parseTreeRoot.entitiesCount())*prob
			elif chosenEntity.children:
				# beta channel
				prob = np.exp(chosenEntity.proposeFromChildLogProb())
				prob += smallNumber
				totalProb += channelProb[1]*(1./self.parseTreeRoot.termEntitiesCount())*prob
				
				# gamma channel
				prob = np.exp(chosenEntity.calAttLogProb())
				prob += smallNumber
				totalProb += channelProb[2]*(1./self.parseTreeRoot.entitiesCount())*prob
			elif chosenEntity.node.isEntity:
				prob = self.cubDist.evaluate(chosenEntity.getCuboid(), chosenEntity.node.symbol)
				prob += smallNumber
				totalProb += channelProb[0]*(1./self.parseTreeRoot.termEntitiesCount())*prob
				
				# gamma channel
				prob = np.exp(chosenEntity.calAttLogProb())
				prob += smallNumber
				totalProb += channelProb[2]*(1./self.parseTreeRoot.entitiesCount())*prob
			else:
				# gamma channel
				prob = np.exp(chosenEntity.calAttLogProb())
				prob += smallNumber
				totalProb += channelProb[2]*(1./self.parseTreeRoot.entitiesCount())*prob
				
			return np.log(totalProb)
			
		
# 		channelProb = [8./10, 1./10, 1./10]
		channelProb = [8./10, 1./10, 1./10]
		channel = rcategorical(channelProb)
		
		if channel == 0:  # alpha channel
			chosenEntity = self.parseTreeRoot.randomTermEntity()
			if chosenEntity == self.parseTreeRoot:
				if self.iteration < 300:
					return self.propose(self.iteration)
				
				newBg = self.bgDist.propose()
				
				backwardLogProbability = proposalLogProb(chosenEntity)
				chosenEntity.assignCuboid(newBg)
				forwardLogProbability = proposalLogProb(chosenEntity)

				if self.saveGraph:
					self.drawGraphAlpha(chosenEntity)
					
				return backwardLogProbability - forwardLogProbability, "Alpha"
			else:
				newCuboid = self.cubDist.propose(chosenEntity.node.symbol)
				
				backwardLogProbability = proposalLogProb(chosenEntity)
				chosenEntity.assignCuboid(newCuboid)
				forwardLogProbability = proposalLogProb(chosenEntity)
				
				if self.saveGraph:
					self.drawGraphAlpha(chosenEntity)
					
				return backwardLogProbability - forwardLogProbability, "Alpha"
			
		elif channel == 1:  # beta channel
			chosenEntity = self.parseTreeRoot.randomParentEntity()
			if chosenEntity == None or chosenEntity == self.parseTreeRoot:
				return self.propose(self.iteration)
			else:
				# Randomly choose a child
				if len(chosenEntity.children) == 0:
					print self.parseTreeRoot.parentEntitiesCount()
					print chosenEntity.node.symbol
					print chosenEntity.node.children
					print chosenEntity.childrenCount
					self.printTree(self.parseTreeRoot, 0)
					
				backwardLogProbability = proposalLogProb(chosenEntity)
				child = chosenEntity.children[pm.distributions.rcategorical([1./len(chosenEntity.children)]*len(chosenEntity.children))]
				chosenEntity.proposeFromChild(child)
				forwardLogProbability = proposalLogProb(chosenEntity)
				
				if self.saveGraph:
					self.drawGraphBeta(chosenEntity, child)
					
				return backwardLogProbability - forwardLogProbability, "Beta"
			
		elif channel == 2:  # gamma channel
			chosenEntity = self.parseTreeRoot.randomEntity()
			if chosenEntity == self.parseTreeRoot:
				return self.propose(self.iteration)
			
			backwardLogProbability = proposalLogProb(chosenEntity)
			chosenEntity.proposeFromParent()
			forwardLogProbability = proposalLogProb(chosenEntity)
			
			if self.saveGraph:
				self.drawGraphGamma(chosenEntity)
				
			return backwardLogProbability - forwardLogProbability, "Gamma"
		
	
	def diffusionGroup(self):
		def proposalLogProb(chosenEntity):
			totalProb = 0
			
			prob = self.cubDist.evaluate(chosenEntity.getCuboid(), chosenEntity.node.symbol)
			prob += smallNumber
			totalProb += channelProb[0]*(1./self.parseTreeRoot.parentEntitiesCount())*prob
			
			# gamma channel
			prob = np.exp(chosenEntity.calAttLogProb())
			prob += smallNumber
			totalProb += channelProb[2]*(1./self.parseTreeRoot.parentEntitiesCount())*prob

			return np.log(totalProb)
			
# 		channelProb = [4./10, 3./10, 3./10]
		channelProb = [4./10, 3./10, 3./10]
		channel = rcategorical(channelProb)
		chosenEntity = self.parseTreeRoot.randomParentEntity()
		if chosenEntity == None or chosenEntity == self.parseTreeRoot:
			return self.propose(self.iteration)
		
		if channel == 0:  # alpha channel
			newCuboid = self.cubDist.propose(chosenEntity.node.symbol)
			
			backwardLogProbability = proposalLogProb(chosenEntity)
			chosenEntity.assignGroupCuboid(newCuboid)
			forwardLogProbability = proposalLogProb(chosenEntity)
			
			if self.saveGraph:
				self.drawGraphAlpha(chosenEntity)
				
			return backwardLogProbability - forwardLogProbability, "Alpha(Group)"
			
		elif channel == 1:  # beta channel
			# Randomly choose a child
			child = chosenEntity.children[pm.distributions.rcategorical([1./len(chosenEntity.children)]*len(chosenEntity.children))]

			backwardLogProbability = chosenEntity.proposeGroupFromChildLogProb(child)
			chosenEntity.proposeGroupFromChild(child)
			forwardLogProbability = chosenEntity.proposeGroupFromChildLogProb(child)

			if self.saveGraph:
				self.drawGraphBeta(chosenEntity, child)
				
			return backwardLogProbability - forwardLogProbability, "Beta(Group)"
			
		elif channel == 2:  # gamma channel
			backwardLogProbability = proposalLogProb(chosenEntity)
			chosenEntity.proposeGroupFromParent()
			forwardLogProbability = proposalLogProb(chosenEntity)
			
			if self.saveGraph:
				self.drawGraphGamma(chosenEntity)
				
			return backwardLogProbability - forwardLogProbability, "Gamma(Group)"

		
		
	def jumpAdd(self):
		parentEntity = self.parseTreeRoot.randomNonTermEntity()
		forwardLogProbability = np.log(1./self.parseTreeRoot.nonTermEntitiesCount())
		newChildNode = parentEntity.node.children[rcategorical([1./len(parentEntity.node.children)]*len(parentEntity.node.children))]
		
		newChildCount = parentEntity.childrenCount[newChildNode.symbol]+1
		if len(parentEntity.node.childrenProb[newChildNode.symbol]) > newChildCount \
		   and parentEntity.node.childrenProb[newChildNode.symbol][newChildCount] > 0:
			newChildEntity = Entity(newChildNode, parentEntity)
			newChildEntity.sampleAttribute()
			newChildNode.sample(newChildEntity)
			parentEntity.children.append(newChildEntity)
			parentEntity.childrenCount[newChildNode.symbol] += 1

			forwardLogProbability += newChildEntity.calNormalizedLogPrior()
			backwardLogProbability = np.log(1./(self.parseTreeRoot.entitiesCount()-1))
		else:
			return self.propose(self.iteration)
		
		if self.saveGraph:
			self.drawGraphAdd(newChildEntity)
			
		return backwardLogProbability - forwardLogProbability, "Add"
		
	def jumpDelete(self):
		def printChildrenCount(entity):
			if entity.node.isTerminal:
				return
			print entity.childrenCount
			if entity.children:
				for child in entity.children:
					printChildrenCount(child)
		
		chosenEntity = self.parseTreeRoot.randomEntity()
		if chosenEntity == self.parseTreeRoot:
			return self.propose(self.iteration)
		
		backwardLogProbability = chosenEntity.calNormalizedLogPrior()
		forwardLogProbability = np.log(1./(self.parseTreeRoot.entitiesCount()-1))
		chosenEntity.delEntity()
		backwardLogProbability -= np.log(float(self.parseTreeRoot.nonTermEntitiesCount()))

		if self.saveGraph:
			self.drawGraphDelete()
			
		return backwardLogProbability - forwardLogProbability, "Delete"
	
	def propose(self, iteration):
		self.iteration = iteration
		
# 		method = rcategorical([4./10, 4./10, 1./10, 1./10])
		method = rcategorical([4./10, 4./10, 1./10, 1./10])
		if method == 0:
			return self.diffusion()
		elif method == 1:
			return self.diffusionGroup()
		elif method == 2:
			return self.jumpAdd()
		elif method == 3:
			return self.jumpDelete()

	def updateTemperature(self, temperature, cubDist):
		def updateDist(node, ratio):
			node.distribution.updateVar(ratio)
			if node.children:
				for child in node.children:
					updateDist(child, ratio)
		
# 		if temperature>=1:
# 			ratio = temperature/self.T
# 			updateDist(self.grammarTreeRoot, ratio)
# 			self.T = temperature
		self.cubDist = cubDist
		
	def updateCubDist(self, cubDist):
		self.cubDist = cubDist

	def calLogPrior(self):
		return self.parseTreeRoot.calNormalizedLogPrior()

	def getEntities(self):
		return self.parseTreeRoot.getData()

	def get_entities(self):
		return self.parseTreeRoot.get_entities()
	
	def calPhysicsEnergy(self):
		return 6*calCollisionVolume(self.getEntities())

	def drawGraphAlpha(self, chosenEntity):
		def addEdge(entity, symbolPar):
			if entity.children:
				for child in entity.children:
					symbolChild = child.node.symbol
					
					if symbolChild not in symbols:
						symbols[symbolChild] = 1
					else:
						symbols[symbolChild] += 1
						symbolChild = symbolChild + str(symbols[symbolChild])
					
					if child == chosenEntity:
						nodeChild = pydot.Node(name=symbolChild, style="filled", fillcolor="red")
						graph.add_node(nodeChild)
						nodeData = pydot.Node(name=" ", color="red", shape="square", penwidth=5)
						graph.add_node(nodeData)
						graph.add_edge(pydot.Edge(nodeChild, nodeData, dir="back", color='red', label='Bottom-up detection', penwidth=5))
						
						graph.add_edge(pydot.Edge(symbolPar,symbolChild))
					else:	
						graph.add_edge(pydot.Edge(symbolPar,symbolChild))
						
					addEdge(child, symbolChild)
		
		graph = pydot.Dot(graph_type='digraph',font='verdana')
		symbols = {}
		addEdge(self.parseTreeRoot, self.parseTreeRoot.node.symbol)
		graph.write_png(self.folder+'graph_'+str(self.iteration)+'.png', prog='dot')

	def drawGraphBeta(self, chosenParent, chosenChild):
		def addEdge(entity, symbolPar):
			if entity.children:
				for child in entity.children:
					symbolChild = child.node.symbol
					
					if symbolChild not in symbols:
						symbols[symbolChild] = 1
					else:
						symbols[symbolChild] += 1
						symbolChild = symbolChild + str(symbols[symbolChild])
						
					if child == chosenParent:
						nodeChosenPar = pydot.Node(name=symbolChild, style="filled", fillcolor="blue")
						graph.add_node(nodeChosenPar)
						graph.add_edge(pydot.Edge(symbolPar,nodeChosenPar))
					elif child == chosenChild:
						graph.add_edge(pydot.Edge(symbolPar, symbolChild, dir="back", color='blue', label='Bottom-up prediction', penwidth=5))
					else:
						graph.add_edge(pydot.Edge(symbolPar,symbolChild))
						
					addEdge(child, symbolChild)
					
		
		graph = pydot.Dot(graph_type='digraph',font='verdana')
		symbols = {}
		addEdge(self.parseTreeRoot, self.parseTreeRoot.node.symbol)
		graph.write_png(self.folder+'graph_'+str(self.iteration)+'.png', prog='dot')

	def drawGraphGamma(self, chosenEntity):
		def addEdge(entity, symbolPar):
			if entity.children:
				for child in entity.children:
					symbolChild = child.node.symbol
					
					if symbolChild not in symbols:
						symbols[symbolChild] = 1
					else:
						symbols[symbolChild] += 1
						symbolChild = symbolChild + str(symbols[symbolChild])
					
					if child == chosenEntity:
						nodeChild = pydot.Node(name=symbolChild, style="filled", fillcolor="green")
						graph.add_node(nodeChild)
						graph.add_edge(pydot.Edge(symbolPar,nodeChild, color='green', label='Top-down prediction', penwidth=5))
					else:
						graph.add_edge(pydot.Edge(symbolPar,symbolChild))
					
					addEdge(child, symbolChild)
		
		graph = pydot.Dot(graph_type='digraph',font='verdana')
		symbols = {}
		addEdge(self.parseTreeRoot, self.parseTreeRoot.node.symbol)
		graph.write_png(self.folder+'graph_'+str(self.iteration)+'.png', prog='dot')

	def drawGraphAdd(self, chosenEntity):
		def addEdge(entity, symbolPar):
			if entity.children:
				for child in entity.children:
					symbolChild = child.node.symbol
					
					if symbolChild not in symbols:
						symbols[symbolChild] = 1
					else:
						symbols[symbolChild] += 1
						symbolChild = symbolChild + str(symbols[symbolChild])
					
					if child == chosenEntity:
						nodeChild = pydot.Node(name=symbolChild, style="filled", fillcolor="green")
						graph.add_node(nodeChild)
						graph.add_edge(pydot.Edge(symbolPar,nodeChild, color='green', label='Top-down', penwidth=5))
					else:
						graph.add_edge(pydot.Edge(symbolPar,symbolChild))
					
					addEdge(child, symbolChild)
		
		graph = pydot.Dot(graph_type='digraph',font='verdana')
		symbols = {}
		addEdge(self.parseTreeRoot, self.parseTreeRoot.node.symbol)
		graph.write_png(self.folder+'graph_'+str(self.iteration)+'.png', prog='dot')

	def drawGraphDelete(self):
		def addEdge(entity, symbolPar):
			if entity.children:
				for child in entity.children:
					symbolChild = child.node.symbol
					
					if symbolChild not in symbols:
						symbols[symbolChild] = 1
					else:
						symbols[symbolChild] += 1
						symbolChild = symbolChild + str(symbols[symbolChild])
					
					graph.add_edge(pydot.Edge(symbolPar,symbolChild))
					addEdge(child, symbolChild)
		
		graph = pydot.Dot(graph_type='digraph',font='verdana')
		symbols = {}
		addEdge(self.parseTreeRoot, self.parseTreeRoot.node.symbol)
		graph.write_png(self.folder+'graph_'+str(self.iteration)+'.png', prog='dot')

	def drawGraph(self, name):
		def addEdge(entity, symbolPar):
			if entity.children:
				for child in entity.children:
					symbolChild = child.node.symbol
					
					if symbolChild not in symbols:
						symbols[symbolChild] = 1
					else:
						symbols[symbolChild] += 1
						symbolChild = symbolChild + str(symbols[symbolChild])
					
					graph.add_edge(pydot.Edge(symbolPar,symbolChild))
					addEdge(child, symbolChild)
		
		graph = pydot.Dot(graph_type='digraph',font='verdana')
		symbols = {}
		addEdge(self.parseTreeRoot, self.parseTreeRoot.node.symbol)
		graph.write_png(os.path.relpath(self.folder+'../' + name), prog='dot')

	def printTree(self, pnode, n):
			print '  '*n + pnode.node.symbol
			if pnode.children:
				for child in pnode.children:
					self.printTree(child, n+1)

	def printNodeTree(self):
		def printNodeTreeRecur(node):
			if node.parent:
				print node.symbol, node, ', parent: '+node.parent.symbol, node.parent
			else:
				print node.symbol, node, ', parent: None'
		
			if node.children:
				for child in node.children:
					printNodeTreeRecur(child)

		printNodeTreeRecur(self.grammarTreeRoot)

	def printParseTree(self):
		def printParseTreeRecur(entity):
			if entity.parent:
				print entity.node.symbol, entity, entity.node.parent.symbol, entity.parent
			else:
				print entity.node.symbol, entity
		
			if entity.children:
				for child in entity.children:
					printParseTreeRecur(child)
					
		printParseTreeRecur(self.parseTreeRoot)

	def writeEntitiesToFile(self, filename):
		outputFile = open(filename, 'w')
		self.parseTreeRoot.writeToFile(outputFile)
		outputFile.close()

